/*
 * Copyright (C) 2021 Baidu, Inc. All Rights Reserved.
 */
package com.chukun.rpc.registry;

import java.io.IOException;
import java.util.Collection;
import java.util.List;

import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.curator.x.discovery.ServiceDiscovery;
import org.apache.curator.x.discovery.ServiceDiscoveryBuilder;
import org.apache.curator.x.discovery.ServiceInstance;
import org.apache.curator.x.discovery.details.JsonInstanceSerializer;

import com.chukun.rpc.core.TinyRpcServiceHelper;
import com.chukun.rpc.core.TinyRpcServiceMetaData;
import com.chukun.rpc.registry.lb.ZKConsistentHashLoadBalancer;

/**
 * <p>zk注册中心实现<p/>
 *
 * @author chukun
 * @version 1.0
 * @since 1.0-SNAPSHOT
 */
public class ZookeeperRegistryService implements RegistryService{
    public static final int BASE_SLEEP_TIME_MS = 1000;
    public static final int MAX_RETRIES = 3;
    public static final String ZK_BASE_PATH = "/tiny_rpc";

    private final ServiceDiscovery<TinyRpcServiceMetaData> serviceDiscovery;


    public ZookeeperRegistryService (String registryAddr){
        CuratorFramework client = CuratorFrameworkFactory.newClient(registryAddr,
                new ExponentialBackoffRetry(BASE_SLEEP_TIME_MS, MAX_RETRIES));
        client.start();

        JsonInstanceSerializer<TinyRpcServiceMetaData> serializer =
                new JsonInstanceSerializer<>(TinyRpcServiceMetaData.class);

        this.serviceDiscovery = ServiceDiscoveryBuilder.builder(TinyRpcServiceMetaData.class)
                .client(client)
                .basePath(ZK_BASE_PATH)
                .serializer(serializer)
                .build();
        try {
            this.serviceDiscovery.start();
        }catch (Exception e) {
            // ignore
            e.printStackTrace();
        }
    }

    @Override
    public void register(TinyRpcServiceMetaData metaData) throws Exception {
        ServiceInstance<TinyRpcServiceMetaData> serviceInstance = ServiceInstance.<TinyRpcServiceMetaData>builder()
                .name(TinyRpcServiceHelper.buildServiceKey(metaData.getServiceName(), metaData.getVersion()))
                .address(metaData.getServiceAddr())
                .port(metaData.getPort())
                .payload(metaData)
                .build();
        this.serviceDiscovery.registerService(serviceInstance);
    }

    @Override
    public void unRegister(TinyRpcServiceMetaData metaData) throws Exception {
        ServiceInstance<TinyRpcServiceMetaData> serviceInstance = ServiceInstance.<TinyRpcServiceMetaData>builder()
                .name(TinyRpcServiceHelper.buildServiceKey(metaData.getServiceName(), metaData.getVersion()))
                .address(metaData.getServiceAddr())
                .port(metaData.getPort())
                .payload(metaData)
                .build();
        this.serviceDiscovery.unregisterService(serviceInstance);
    }

    @Override
    public TinyRpcServiceMetaData discovery(String serviceName, int hash) throws Exception {
        Collection<ServiceInstance<TinyRpcServiceMetaData>> serviceInstances =
                this.serviceDiscovery.queryForInstances(serviceName);
        if (serviceInstances.isEmpty()) {
            return null;
        }
        ZKConsistentHashLoadBalancer loadBalancer = new ZKConsistentHashLoadBalancer();
        ServiceInstance<TinyRpcServiceMetaData> instance =
                loadBalancer.select((List<ServiceInstance<TinyRpcServiceMetaData>>) serviceInstances, hash);
        if (instance == null) {
            return  serviceInstances.iterator().next().getPayload();
        }
        return instance.getPayload();
    }

    @Override
    public void destroy() {
        try {
            serviceDiscovery.close();
        }catch (IOException e) {
            // ignore
            e.printStackTrace();
        }
    }
}
